import tushare as ts
import numpy as np
import matplotlib.pyplot as plt
from pandas import read_csv
from pandas import DataFrame
from keras.layers.core import Flatten
from keras.models import Sequential
from keras.layers import Dense, Dropout
from keras.layers import LSTM
from keras.layers import TimeDistributed


def data_normalize(df):
    norm = df.apply(lambda x: (x - np.min(x)) / (np.max(x) - np.min(x)))
    return norm


def data_split(df, ref_day=5, predict_day=1):

    X_train, Y_train = [], []

    for i in range(df.shape[0] - predict_day - ref_day):

        X_train.append(np.array(df.iloc[i:i+ref_day, :]))
        Y_train.append(
            np.array(df.iloc[i+ref_day:i+ref_day+predict_day]["close"]))

    pos = int(len(Y_train) * 0.75)
    return np.array(X_train)[:pos], np.array(X_train)[pos + 1:], np.array(Y_train)[:pos], np.array(Y_train)[pos + 1:]


def lstm_stock_model():

    model = Sequential()

    model.add(LSTM(64, input_shape=(5, 5),
                   activation=None, return_sequences=False))
    model.add(Dropout(0.5))
    #model.add(Flatten())
    model.add(Dense(1, activation='relu'))

    model.compile(loss="mae", optimizer="adam",
                  metrics=['mae'])

    model.summary()
    return model


def main():
    
    data = ts.get_k_data('600036', ktype='D', start='2010-01-01', end='2019-12-20')
    data.to_csv('kdata.csv')
    
    #data = read_csv('kdata.csv')

    xaxis = range(1, len(data['date']), 240)
    xaxisnew = range(2010, 2021)

    '''
    plt.subplot(511)
    plt.plot(data['date'], data['open'])
    plt.xticks(xaxis, xaxisnew)
    plt.ylabel('open')
    plt.subplot(512)
    plt.plot(data['date'], data['close'])
    plt.xticks(xaxis, xaxisnew)
    plt.ylabel('close')
    plt.subplot(513)
    plt.plot(data['date'], data['high'])
    plt.xticks(xaxis, xaxisnew)
    plt.ylabel('high')
    plt.subplot(514)
    plt.plot(data['date'], data['low'])
    plt.xticks(xaxis, xaxisnew)
    plt.ylabel('low')
    plt.subplot(515)
    plt.plot(data['date'], data['volume'])
    plt.xticks(xaxis, xaxisnew)
    plt.ylabel('volume')
    plt.show()
    '''

    data = data[['open', 'close', 'high', 'low', 'volume']]
    df = data_normalize(data)

    X_train, X_test, y_train, y_test = data_split(df)

    model = lstm_stock_model()

    his = model.fit(X_train, y_train, epochs=100, batch_size=100)
    plt.plot(his.history['loss'])
    plt.show()

    res = model.predict(X_test)

    close = np.array(data['close'])
    res = res * (np.max(close) - np.min(close)) + np.min(close)

    r = model.predict(X_train)
    rr = r * (np.max(close) - np.min(close)) + np.min(close)
    plt.plot(close, 'r-', range(len(y_train)), rr, 'b-')

    plt.plot(range(len(close) - len(y_test), len(close)), res, 'g-')
    plt.xticks(xaxis, xaxisnew)
    plt.ylabel('close')
    plt.show()


if __name__ == '__main__':
    main()
